from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
from datetime import datetime
import json
import os

app = FastAPI(title="AI Training Log System API")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 在生产环境中应该限制为前端域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 日志数据模型
class LogEntry(BaseModel):
    timestamp: str
    level: str
    message: str
    source: Optional[str] = None

# 训练日志响应模型
class TrainLogResponse(BaseModel):
    code: int
    message: str
    data: Dict[str, Any]

# 模拟日志数据
sample_logs = [
    {
        "timestamp": datetime.now().isoformat(),
        "level": "INFO",
        "message": "应用启动成功",
        "source": "system"
    },
    {
        "timestamp": datetime.now().isoformat(),
        "level": "WARNING",
        "message": "磁盘空间不足",
        "source": "system"
    },
    {
        "timestamp": datetime.now().isoformat(),
        "level": "ERROR",
        "message": "数据库连接失败",
        "source": "database"
    }
]

# 日志文件路径
LOG_FILE = "logs.json"
# 训练日志目录
TRAIN_LOGS_DIR = "train_logs"

# 确保训练日志目录存在
os.makedirs(TRAIN_LOGS_DIR, exist_ok=True)

# 如果日志文件存在，则加载；否则使用示例数据
def load_logs():
    if os.path.exists(LOG_FILE):
        try:
            with open(LOG_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"Error loading logs: {e}")
            return sample_logs
    else:
        # 保存示例数据到文件
        save_logs(sample_logs)
        return sample_logs

# 保存日志到文件
def save_logs(logs):
    with open(LOG_FILE, "w", encoding="utf-8") as f:
        json.dump(logs, f, ensure_ascii=False, indent=2)

# 创建示例训练日志文件（用于演示）
def create_sample_train_log(task_code, train_code):
    task_dir = os.path.join(TRAIN_LOGS_DIR, task_code)
    os.makedirs(task_dir, exist_ok=True)
    
    train_dir = os.path.join(task_dir, train_code)
    os.makedirs(train_dir, exist_ok=True)
    
    log_file = os.path.join(train_dir, "train.log")
    
    # 如果文件不存在，创建示例日志
    if not os.path.exists(log_file):
        with open(log_file, "w", encoding="utf-8") as f:
            f.write("2023-05-18 10:00:00 [INFO] 开始训练任务\n")
            f.write("2023-05-18 10:00:05 [INFO] 加载数据集: 10000个样本\n")
            f.write("2023-05-18 10:00:10 [INFO] 初始化模型参数\n")
            f.write("2023-05-18 10:00:15 [INFO] Epoch 1/100 开始\n")
            f.write("2023-05-18 10:01:00 [INFO] Epoch 1/100 完成, loss: 2.345, accuracy: 0.45\n")
            f.write("2023-05-18 10:01:05 [INFO] Epoch 2/100 开始\n")
            f.write("2023-05-18 10:02:00 [INFO] Epoch 2/100 完成, loss: 1.987, accuracy: 0.52\n")
    
    return log_file

@app.get("/")
def read_root():
    return {"message": "Welcome to AI Training Log System API"}

@app.get("/api/logs", response_model=List[LogEntry])
def get_logs():
    return load_logs()

@app.post("/api/logs", response_model=LogEntry)
def add_log(log: LogEntry):
    logs = load_logs()
    logs.append(log.dict())
    save_logs(logs)
    return log

@app.get("/api/logs/level/{level}", response_model=List[LogEntry])
def get_logs_by_level(level: str):
    logs = load_logs()
    filtered_logs = [log for log in logs if log["level"].upper() == level.upper()]
    return filtered_logs

@app.get("/train/getTrainLog", response_model=TrainLogResponse)
def api_get_train_log(
    task_code: str = Query(..., description="任务代码"),
    train_code: str = Query(..., description="训练代码"),
    train_log_index: int = Query(0, description="已获取的日志行数")
):
    # 检查任务是否存在
    task_dir = os.path.join(TRAIN_LOGS_DIR, task_code)
    if not os.path.exists(task_dir):
        # 为了演示，创建示例日志
        create_sample_train_log(task_code, train_code)
    
    # 检查训练是否存在
    train_dir = os.path.join(task_dir, train_code)
    if not os.path.exists(train_dir):
        # 为了演示，创建示例日志
        create_sample_train_log(task_code, train_code)
    
    # 检查日志文件是否存在
    log_file = os.path.join(train_dir, "train.log")
    if not os.path.exists(log_file):
        return TrainLogResponse(
            code=404,
            message="训练日志文件不存在",
            data={"log_lines": [], "train_log_index": train_log_index}
        )
    
    # 读取日志文件
    try:
        with open(log_file, "r", encoding="utf-8") as f:
            all_lines = f.readlines()
        
        # 获取新的日志行
        new_lines = all_lines[train_log_index:]
        new_train_log_index = len(all_lines)
        
        return TrainLogResponse(
            code=200,
            message="成功获取训练日志",
            data={
                "log_lines": new_lines,
                "train_log_index": new_train_log_index
            }
        )
    except Exception as e:
        return TrainLogResponse(
            code=500,
            message=f"读取日志文件失败: {str(e)}",
            data={"log_lines": [], "train_log_index": train_log_index}
        )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)